import pandas as pd
import numpy as np
from itertools import combinations
import xlrd

def preprocess_dataframe(df):
    # 1. Remove spaces from 分子式
    df['分子式'] = df['分子式'].str.strip()
    
    # 2. Round m/z and 质量 to 2 decimal places
    # df['m/z'] = df['m/z'].round(2)
    # df['质量'] = df['质量'].round(2)
    
    # 3. Extract 部位 and 极性 from 文件
    df[['部位', '极性']] = df['文件'].str.extract(r'TZS-([^-]+)-([^.]+)\.d')
    
    # 4. Create m/z-质量
    df['差'] = df['m/z']-df['质量']
    
    return df


# print(processed_df[['部位','极性','差']])


def create_new_dataframe(df):
    # Create new dataframe with specified columns
    new_df = pd.DataFrame()
    
    # 1. Empty "No." column
    new_df['No.'] = ''
    
    # 2-7. Copy specified columns
    new_df['部位'] = df['部位']
    new_df['Mode'] = df['极性']
    new_df['Name'] = df['名称']
    new_df['Formula'] = df['分子式'].str.replace(' ', '')
    new_df['RT'] = df['RT']
    new_df['m/z'] = df['m/z']
    new_df['Mass'] = df['质量']
    
    # 8. Create Adduct column based on "差" values
    def determine_adduct(x):
        if 0.95 <= x <= 1.05:
            return '[M+H]+'
        elif -1.05 <= x <= -0.95:
            return '[M-H]-'
        elif 44.95 <= x <= 45.05:
            return '[M+COOH]-'
        elif 22.95 <= x <= 23.05:
            return '[M+NA]+'
        else:
            return ''
    new_df['Adduct'] = df['差'].apply(determine_adduct)
    
    # 9-12. Copy specified columns
    new_df['Mass error'] = df['误差(Tgt, ppm)']
    new_df['Score'] = df['分数 (Tgt)']
    new_df['Height'] = df['峰高']
    new_df['Area'] = df['面积']
    
    # 13-15. Empty columns
    new_df['峰面积百分比'] = ''
    new_df['CID'] = ''
    new_df['CAS'] = ''
    
    # 16. Create "正负对应" column
    # First, round Mass to 2 decimals for comparison
    df['Mass_rounded'] = df['质量'].round(2)
    
    # Initialize column
    new_df['正负对应'] = ''
    
    # Group counter
    group_counter = 1
    
    # Get unique 部位 values
    unique_buwei = df['部位'].unique()
    
    # For each 部位
    for buwei in unique_buwei:
        # Get rows for this 部位
        buwei_data = df[df['部位'] == buwei]
        
        # Get unique polarities for this 部位
        polarities = buwei_data['极性'].unique()
        
        # If we have both polarities
        if len(polarities) > 1:
            # Get rows for each polarity
            for pol1, pol2 in combinations(polarities, 2):
                data1 = buwei_data[buwei_data['极性'] == pol1]
                data2 = buwei_data[buwei_data['极性'] == pol2]
                
                # Compare each pair
                for _, row1 in data1.iterrows():
                    for _, row2 in data2.iterrows():
                        if (row1['Mass_rounded'] == row2['Mass_rounded'] and 
                            abs(row1['RT'] - row2['RT']) < 0.5):
                            # Mark both rows
                            new_df.loc[new_df.index == row1.name, '正负对应'] = str(group_counter)
                            new_df.loc[new_df.index == row2.name, '正负对应'] = str(group_counter)
                            group_counter += 1
    
    # 17. Create "同时出峰" column
    new_df['同时出峰'] = ''
    group_counter = 1
    
    # For each polarity
    for polarity in df['极性'].unique():
        polarity_data = df[df['极性'] == polarity]
        
        # Get unique 部位 values for this polarity
        buwei_values = polarity_data['部位'].unique()
        
        # If we have multiple 部位 values
        if len(buwei_values) > 1:
            # Compare each pair of different 部位
            for buwei1, buwei2 in combinations(buwei_values, 2):
                data1 = polarity_data[polarity_data['部位'] == buwei1]
                data2 = polarity_data[polarity_data['部位'] == buwei2]
                
                # Compare each pair
                for _, row1 in data1.iterrows():
                    for _, row2 in data2.iterrows():
                        if (row1['Mass_rounded'] == row2['Mass_rounded'] and 
                            abs(row1['RT'] - row2['RT']) < 0.5):
                            # Mark both rows
                            new_df.loc[new_df.index == row1.name, '同时出峰'] = str(group_counter)
                            new_df.loc[new_df.index == row2.name, '同时出峰'] = str(group_counter)
                            group_counter += 1
    
    return new_df

# Usage:
# Read the original xls file
# df = pd.read_excel('your_file.xls', engine='xlrd')

# new_df.to_excel('processed_file.xlsx', index=False)

def filter_low_scores(df):
    
    # Create a copy
    filtered_df = df.copy()
    
    # Filter satisfied conditions
    filtered_df = filtered_df[
        ((filtered_df['正负对应'] != '') |
         (filtered_df['同时出峰'] != '') |
         (filtered_df['Score'] >= 75))
    ] 
    
    filtered_df = filtered_df.reset_index(drop=True)
    
    return filtered_df


def round_decimals(df):
    
    # Round decimals
    df['m/z'] = df['m/z'].round(2)
    df['Mass'] = df['Mass'].round(2)
    
    return df

df = pd.read_excel('太子参示例.xls', engine='xlrd')

processed_df = preprocess_dataframe(df)

new_df = create_new_dataframe(processed_df)

filtered_df = filter_low_scores(new_df)

rounded_df = round_decimals(filtered_df)

percent_df = rounded_df.copy()                    
area_sums = percent_df.groupby(['部位', 'Mode'])['Area'].transform('sum')
percent_df['峰面积百分比'] = (percent_df['Area'] / area_sums * 100).round(2)
 
percent_df.to_excel('processed_file.xlsx', index=False)                                         